# @time     :2022/3/6
# @Author   :Simple
# @FileName :auto_make_api_yaml
import os
import sys
from urllib.parse import urlparse, parse_qs, unquote_plus

from mitmproxy import http

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from util.config import api_url, BASE_DIR, api_yaml_template, init_dir, API_DIR


def handle_url(url):
    """
    获取url中的query的key
    :param url:
    :return:
    """
    # url解码 或者:unquote(url)
    url = unquote_plus(url)
    # 分解url，并获取query中的key
    parsed = urlparse(url)
    param_keys = parse_qs(parsed.query).keys()
    return param_keys


class AutoGetApi:
    """
    抓取request，并且自动生成对应的api yaml文件
    使用方法：
        1、开启服务和端口（默认8080）：mitmdump -s util/auto_make_api_yaml.py
        2、设备设置代理
        3、设备发送请求
        4、mitmdump自动抓取request并且生成对应yaml文件
    """
    def request(self, flow: http.HTTPFlow):
        """
        抓取request，解析出url、method、query、data，并且按模板写api
        :param flow:
        :return:
        """
        url = flow.request.pretty_url
        # 抓取请求，url包含api_url
        if api_url in url:
            host_url = url.split("?")[0]
            parsed = urlparse(url)
            path_url = parsed.path
            path_dir = os.path.split(path_url)[0]
            yaml_name = os.path.split(path_url)[1]+'.yml'
            method = flow.request.method
            # head = flow.request.headers
            params_keys = handle_url(url)
            data_keys = flow.request.json().keys()
            print(f"method:{method},  url:{host_url}, path_dir:{path_dir},yaml_name:{yaml_name}, param_keys:{params_keys}, data_json:{data_keys}")

            # 读取模板文件，并拼接参数
            with open(api_yaml_template) as f:
                data = f.read()
                new_data = data.format(method=method, url=host_url)

            # 写api的yaml文件
            path_dir_list = path_dir[1:].split("/")
            path_dir = API_DIR
            for p in path_dir_list:
                path_dir = init_dir(path_dir, p)
            yaml_path = os.path.join(path_dir, yaml_name)
            with open(yaml_path, "w", encoding="utf-8") as f:
                f.write(new_data)
                if len(params_keys) > 0:
                    f.write('"params": \n')
                    for key in params_keys:
                        f.write(f'  "{key}": ${key}\n')
                if len(data_keys) > 0:
                    f.write('"json": \n')
                    for key in data_keys:
                        f.write(f'  "{key}": ${key}\n')

    def response(self, flow: http.HTTPFlow):
        pass


addons = [
    AutoGetApi()
]

